#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文生图服务

使用火山引擎方舟平台的文生图API，提供文本到图像的生成服务
"""

import os
import sys
from typing import List, Optional, Dict, Any
from volcenginesdkarkruntime import Ark
from volcenginesdkarkruntime.types.images.images import SequentialImageGenerationOptions


class TextToImageService:
    """
    文生图服务类
    用于将文本描述转换为图像
    """
    
    # 最大图片生成数量限制
    MAX_IMAGE_COUNT = 5
    
    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        """
        初始化文生图服务
        
        Args:
            api_key: API密钥，如果不提供则从环境变量或配置中读取
            base_url: API基础URL
        """
        # 配置API密钥
        self.api_key = api_key or os.getenv("ARK_API_KEY") or "a2f51613-131e-4ded-a077-c087b2ea3f4e"
        
        # 配置基础URL
        self.base_url = base_url or "https://ark.cn-beijing.volces.com/api/v3"
        
        # 配置模型
        self.model = "doubao-seedream-4-0-250828"
        
        try:
            # 初始化Ark客户端
            self.client = Ark(
                base_url=self.base_url,
                api_key=self.api_key,
            )
            print("✅ 文生图服务初始化成功!")
            
        except Exception as e:
            print(f"❌ 文生图服务初始化失败: {e}")
            sys.exit(1)
    
    def generate_images(
        self,
        prompt: str,
        num_images: int = 1,
        size: str = "2K",
        reference_images: Optional[List[str]] = None,
        watermark: bool = True,
        response_format: str = "url"
    ) -> Dict[str, Any]:
        """
        根据文本描述生成图片
        
        Args:
            prompt: 图片生成的文本描述
            num_images: 生成图片数量（最多5张）
            size: 图片尺寸，支持 "1K"、"2K" 等
            reference_images: 参考图片URL列表（可选）
            watermark: 是否添加水印
            response_format: 返回格式，"url" 或 "b64_json"
            
        Returns:
            包含生成结果的字典，格式为:
            {
                "success": True/False,
                "message": "提示信息",
                "data": {
                    "images": [
                        {"url": "图片URL", "size": "图片尺寸"},
                        ...
                    ],
                    "count": 实际生成数量
                }
            }
        """
        try:
            # 验证输入参数
            if not prompt or not prompt.strip():
                return {
                    "success": False,
                    "message": "图片描述不能为空",
                    "data": None
                }
            
            # 限制生成数量
            if num_images < 1:
                num_images = 1
            elif num_images > self.MAX_IMAGE_COUNT:
                num_images = self.MAX_IMAGE_COUNT
                print(f"⚠️  生成数量超过限制，已调整为最大值 {self.MAX_IMAGE_COUNT}")
            
            # 构建请求参数
            generate_params = {
                "model": self.model,
                "prompt": prompt.strip(),
                "size": size,
                "sequential_image_generation": "auto",
                "sequential_image_generation_options": SequentialImageGenerationOptions(max_images=num_images),
                "response_format": response_format,
                "watermark": watermark
            }
            
            # 如果提供了参考图片，添加到参数中
            if reference_images and isinstance(reference_images, list) and len(reference_images) > 0:
                generate_params["image"] = reference_images
            
            print(f"🎨 开始生成图片，描述: {prompt}, 数量: {num_images}")
            
            # 调用API生成图片
            images_response = self.client.images.generate(**generate_params)
            
            # 解析响应
            image_list = []
            for image in images_response.data:
                image_info = {
                    "url": image.url if hasattr(image, 'url') else None,
                    "size": image.size if hasattr(image, 'size') else size
                }
                
                # 如果返回格式是 b64_json，添加 base64 数据
                if response_format == "b64_json" and hasattr(image, 'b64_json'):
                    image_info["b64_json"] = image.b64_json
                
                image_list.append(image_info)
                print(f"✅ 图片生成成功: {image_info.get('url', 'base64数据')}")
            
            return {
                "success": True,
                "message": f"成功生成 {len(image_list)} 张图片",
                "data": {
                    "images": image_list,
                    "count": len(image_list),
                    "prompt": prompt,
                    "size": size
                }
            }
            
        except Exception as e:
            error_msg = f"图片生成失败: {str(e)}"
            print(f"❌ {error_msg}")
            return {
                "success": False,
                "message": error_msg,
                "data": None
            }
    
    def validate_prompt(self, prompt: str) -> tuple[bool, str]:
        """
        验证提示词是否合法
        
        Args:
            prompt: 待验证的提示词
            
        Returns:
            (是否合法, 提示信息)
        """
        if not prompt or not prompt.strip():
            return False, "提示词不能为空"
        
        # 可以在这里添加更多的验证逻辑
        # 例如：长度限制、敏感词检测等
        
        if len(prompt.strip()) < 5:
            return False, "提示词太短，请提供更详细的描述"
        
        if len(prompt.strip()) > 1000:
            return False, "提示词太长，请精简描述"
        
        return True, "验证通过"


# 如果直接运行此文件，进行测试
if __name__ == "__main__":
    # 创建服务实例
    service = TextToImageService()
    
    # 测试生成图片
    result = service.generate_images(
        prompt="生成3张女孩和奶牛玩偶在游乐园开心地坐过山车的图片，涵盖早晨、中午、晚上",
        num_images=3,
        size="2K"
    )
    
    # 输出结果
    print("\n" + "="*50)
    print("生成结果:")
    print(f"成功: {result['success']}")
    print(f"消息: {result['message']}")
    if result['data']:
        print(f"生成数量: {result['data']['count']}")
        for idx, img in enumerate(result['data']['images'], 1):
            print(f"图片 {idx}: {img['url']}")